import os
import json
import zipfile
import argparse
from pathlib import Path
from typing import List, Dict, Any, Optional

class ImageZipProcessor:
    """
    图片文件处理与打包工具类
    
    主要功能：
    1. 处理模式：递归遍历文件夹，为每个图片文件创建JSON并打包成ZIP，然后清理原文件
    2. 读取模式：读取ZIP文件中的JSON信息和文件列表
    3. 更新模式：更新ZIP文件中的JSON内容
    
    工作机制：
    - 处理模式：
        1. 递归遍历指定文件夹及其所有子文件夹
        2. 查找所有的图片文件（.jpg, .jpeg, .png 及其大写版本）
        3. 为每个图片文件创建一个同名的JSON文件，内容为 {"file": "文件名", "pattern": ["父文件夹名"]}
        4. 将图片文件和对应的JSON文件打包成ZIP压缩包，保存在目标文件夹的pyma子文件夹中
        5. 清理原文件：删除原始的图片文件和JSON文件
    
    - 读取模式：
        1. 读取ZIP文件中的JSON文件内容
        2. 获取ZIP文件中包含的所有文件列表
    
    - 更新模式：
        1. 更新ZIP文件中的JSON文件内容
        2. 保持其他文件不变
    """
    
    def __init__(self, verbose: bool = True):
        """
        初始化处理器
        
        Args:
            verbose: 是否显示详细输出
        """
        self.verbose = verbose
        # 支持的图片格式（包括大小写）
        self.supported_formats = [
            '*.jpg', '*.JPG', '*.jpeg', '*.JPEG', 
            '*.png', '*.PNG'
        ]
    
    def _log(self, message: str):
        """日志输出"""
        if self.verbose:
            print(message)
    
    def process_folder(self, root_folder: str, keep_original: bool = False) -> int:
        """
        处理文件夹：为所有图片文件创建JSON并打包，然后清理原文件
        
        Args:
            root_folder: 要处理的根文件夹路径
            keep_original: 是否保留原始文件（图片和JSON）
            
        Returns:
            int: 成功处理的文件数量
        """
        root_path = Path(root_folder)
        processed_count = 0
        
        if not root_path.exists():
            raise FileNotFoundError(f"路径 '{root_folder}' 不存在")
        
        if not root_path.is_dir():
            raise NotADirectoryError(f"'{root_folder}' 不是一个文件夹")
        
        self._log(f"开始处理文件夹: {root_folder}")
        if keep_original:
            self._log("模式: 将保留原始文件")
        else:
            self._log("模式: 将在打包后清理原始文件")
        
        # 创建pyma输出目录
        pyma_dir = root_path / "pyma"
        pyma_dir.mkdir(exist_ok=True)
        self._log(f"输出目录: {pyma_dir}")
        
        # 遍历所有支持的图片格式
        for pattern in self.supported_formats:
            for image_file in root_path.rglob(pattern):
                if self._process_single_image_file(image_file, pyma_dir, keep_original):
                    processed_count += 1
        
        self._log(f"处理完成！成功处理 {processed_count} 个文件")
        return processed_count
    
    def _process_single_image_file(self, image_path: Path, output_dir: Path, keep_original: bool) -> bool:
        """
        处理单个图片文件
        
        Args:
            image_path: 图片文件路径
            output_dir: 输出目录（pyma文件夹）
            keep_original: 是否保留原始文件
            
        Returns:
            bool: 处理成功返回True
        """
        try:
            # 获取文件名和父文件夹名
            filename = image_path.name
            parent_dir_name = image_path.parent.name
            
            # 创建JSON文件
            json_path = image_path.with_suffix('.json')
            json_content = {
                "file": filename,
                "pattern": [parent_dir_name]
            }
            
            with open(json_path, 'w', encoding='utf-8') as json_file:
                json.dump(json_content, json_file, ensure_ascii=False, indent=2)
            
            self._log(f"✓ 创建JSON文件: {json_path}")
            
            # 创建ZIP文件（保存在pyma文件夹中）
            zip_filename = output_dir / f"{image_path.stem}.zip"
            success = self._create_zip_archive(zip_filename, image_path, json_path)
            
            # 清理原始文件（如果不需要保留）
            if not keep_original and success:
                self._cleanup_original_files(image_path, json_path)
            
            return success
            
        except Exception as e:
            self._log(f"✗ 处理文件 {image_path} 时出错: {e}")
            return False
    
    def _create_zip_archive(self, zip_path: Path, image_path: Path, json_path: Path) -> bool:
        """
        创建ZIP压缩包
        
        Args:
            zip_path: ZIP文件路径（在pyma文件夹中）
            image_path: 图片文件路径
            json_path: JSON文件路径
            
        Returns:
            bool: 创建成功返回True
        """
        try:
            with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
                zipf.write(image_path, image_path.name)
                self._log(f"  ↳ 添加图片到ZIP: {image_path.name}")
                
                zipf.write(json_path, json_path.name)
                self._log(f"  ↳ 添加JSON到ZIP: {json_path.name}")
            
            self._log(f"✓ 创建ZIP文件: {zip_path}")
            return True
            
        except Exception as e:
            self._log(f"✗ 创建ZIP文件 {zip_path} 时出错: {e}")
            return False
    
    def _cleanup_original_files(self, image_path: Path, json_path: Path):
        """
        清理原始文件
        
        Args:
            image_path: 原始图片文件路径
            json_path: JSON文件路径
        """
        try:
            # 删除JSON文件
            if json_path.exists():
                json_path.unlink()
                self._log(f"✓ 清理JSON文件: {json_path}")
            
            # 删除原始图片文件
            if image_path.exists():
                image_path.unlink()
                self._log(f"✓ 清理原始图片文件: {image_path}")
                
        except Exception as e:
            self._log(f"✗ 清理文件时出错: {e}")
    
    def get_json_info(self, zip_path: str) -> Optional[Dict[str, Any]]:
        """
        获取ZIP文件中的JSON信息
        
        Args:
            zip_path: ZIP文件路径
            
        Returns:
            Optional[Dict]: JSON内容字典，如果失败返回None
        """
        try:
            with zipfile.ZipFile(zip_path, 'r') as zipf:
                # 查找JSON文件
                json_files = [f for f in zipf.namelist() if f.lower().endswith('.json')]
                if not json_files:
                    self._log(f"✗ 在 {zip_path} 中未找到JSON文件")
                    return None
                
                # 读取第一个JSON文件
                json_filename = json_files[0]
                with zipf.open(json_filename) as json_file:
                    json_content = json.load(json_file)
                    self._log(f"✓ 读取JSON信息: {zip_path}")
                    return json_content
                    
        except Exception as e:
            self._log(f"✗ 读取ZIP文件 {zip_path} 时出错: {e}")
            return None
    
    def get_zip_contents(self, zip_path: str) -> Optional[List[str]]:
        """
        获取ZIP文件中包含的所有文件列表
        
        Args:
            zip_path: ZIP文件路径
            
        Returns:
            Optional[List]: 文件列表，如果失败返回None
        """
        try:
            with zipfile.ZipFile(zip_path, 'r') as zipf:
                file_list = zipf.namelist()
                self._log(f"✓ 获取ZIP内容: {zip_path} 包含 {len(file_list)} 个文件")
                return file_list
                
        except Exception as e:
            self._log(f"✗ 读取ZIP文件 {zip_path} 时出错: {e}")
            return None
    
    def update_json_info(self, zip_path: str, new_pattern: List[str]) -> bool:
        """
        更新ZIP文件中的JSON信息
        
        Args:
            zip_path: ZIP文件路径
            new_pattern: 新的pattern列表
            
        Returns:
            bool: 更新成功返回True
        """
        temp_dir = None
        try:
            # 创建临时目录
            temp_dir = Path("temp_zip_update")
            temp_dir.mkdir(exist_ok=True)
            
            # 读取原始ZIP文件
            with zipfile.ZipFile(zip_path, 'r') as zipf:
                # 获取所有文件列表
                file_list = zipf.namelist()
                
                # 查找JSON文件
                json_files = [f for f in file_list if f.lower().endswith('.json')]
                if not json_files:
                    self._log(f"✗ 在 {zip_path} 中未找到JSON文件")
                    return False
                
                json_filename = json_files[0]
                
                # 提取所有文件到临时目录
                zipf.extractall(temp_dir)
            
            # 更新JSON文件（保持filename不变，只更新pattern）
            json_file_path = temp_dir / json_filename
            with open(json_file_path, 'r', encoding='utf-8') as json_file:
                existing_content = json.load(json_file)
            
            # 更新pattern，保持原有的filename
            new_json_content = {
                "file": existing_content.get("file", ""),
                "pattern": new_pattern
            }
            
            with open(json_file_path, 'w', encoding='utf-8') as json_file:
                json.dump(new_json_content, json_file, ensure_ascii=False, indent=2)
            
            # 重新创建ZIP文件
            with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
                for file_path in temp_dir.rglob('*'):
                    if file_path.is_file():
                        arcname = file_path.relative_to(temp_dir)
                        zipf.write(file_path, arcname)
            
            self._log(f"✓ 更新JSON信息: {zip_path} -> {new_pattern}")
            return True
            
        except Exception as e:
            self._log(f"✗ 更新ZIP文件 {zip_path} 时出错: {e}")
            return False
        finally:
            # 清理临时目录
            if temp_dir and temp_dir.exists():
                import shutil
                shutil.rmtree(temp_dir)
    
    def dry_run(self, folder_path: str, keep_original: bool = False):
        """
        模拟运行，显示将要执行的操作
        
        Args:
            folder_path: 要模拟处理的文件夹路径
            keep_original: 是否显示保留原始文件
        """
        self._log("模拟运行模式 - 显示将要执行的操作：")
        self._log("=" * 60)
        
        root_path = Path(folder_path)
        file_count = 0
        
        for pattern in self.supported_formats:
            for image_file in root_path.rglob(pattern):
                filename = image_file.name
                parent_dir = image_file.parent.name
                json_file = image_file.with_suffix('.json')
                zip_file = root_path / "pyma" / f"{image_file.stem}.zip"
                
                self._log(f"图片文件: {image_file}")
                self._log(f"  JSON文件: {json_file}")
                self._log(f"  JSON内容: {{'file': '{filename}', 'pattern': ['{parent_dir}']}}")
                self._log(f"  ZIP文件: {zip_file}")
                self._log(f"  ZIP内容: {image_file.name}, {json_file.name}")
                
                if keep_original:
                    self._log(f"  文件保留: 原始图片和JSON文件将保留")
                else:
                    self._log(f"  文件清理: 原始图片和JSON文件将被删除")
                
                self._log("-" * 40)
                
                file_count += 1
        
        self._log(f"总计发现 {file_count} 个图片文件需要处理")


def setup_argparse():
    """
    设置命令行参数解析
    """
    parser = argparse.ArgumentParser(
        description="""
图片文件处理与打包工具

工作机制：

1. 处理模式（默认）：
   - 递归遍历指定文件夹及其所有子文件夹
   - 查找所有的图片文件（支持：.jpg, .jpeg, .png 及其大写版本）
   - 为每个图片文件创建一个同名的JSON文件，内容为 {"file": "文件名", "pattern": ["父文件夹名"]}
   - 将图片文件和对应的JSON文件打包成ZIP压缩包，保存在目标文件夹的pyma子文件夹中
   - 清理原文件：删除原始的图片文件和JSON文件（默认行为）

2. 读取模式（--get-json 或 --get-contents）：
   - 读取ZIP文件中的JSON文件内容或文件列表
   - 支持单个文件或文件夹批量处理

3. 更新模式（--update-json）：
   - 更新ZIP文件中的JSON文件内容
   - 保持其他文件不变

文件处理说明：
  - 支持格式：JPG, JPEG, PNG（大小写不敏感）
  - JSON格式：{"file": "文件名", "pattern": ["文件夹名"]}
  - 输出位置：在目标文件夹下创建pyma子文件夹存放所有ZIP文件
  - 默认行为：创建ZIP文件后，自动删除原始图片文件和临时JSON文件
  - 保留文件：使用 --keep-original 参数可以保留原始文件

示例文件结构：
  输入： photos/nature/forest.jpg
  处理过程：
    - 创建 photos/nature/forest.json (内容: {"file": "forest.jpg", "pattern": ["nature"]})
    - 创建 photos/pyma/forest.zip (包含forest.jpg和forest.json)
    - 删除 photos/nature/forest.jpg 和 photos/nature/forest.json
  最终结果：只有 photos/pyma/forest.zip 保留
""",
        formatter_class=argparse.RawDescriptionHelpFormatter
    )
    
    # 主要操作模式
    mode_group = parser.add_argument_group('操作模式')
    mode_group.add_argument(
        'folder',
        nargs='?',
        help='要处理的文件夹路径（处理模式使用）'
    )
    mode_group.add_argument(
        '--get-json',
        metavar='ZIP_FILE',
        help='读取ZIP文件中的JSON信息'
    )
    mode_group.add_argument(
        '--get-contents',
        metavar='ZIP_FILE',
        help='获取ZIP文件中包含的文件列表'
    )
    mode_group.add_argument(
        '--update-json',
        metavar=('ZIP_FILE', 'PATTERN'),
        nargs=2,
        help='更新ZIP文件中的JSON信息，例如: --update-json file.zip "new_pattern1,new_pattern2"'
    )
    
    # 选项参数
    option_group = parser.add_argument_group('选项参数')
    option_group.add_argument(
        '--keep-original',
        action='store_true',
        help='在处理完成后保留原始文件（图片和JSON），默认会清理原始文件'
    )
    option_group.add_argument(
        '--dry-run',
        action='store_true',
        help='模拟运行，显示将要执行的操作但不实际创建文件'
    )
    option_group.add_argument(
        '--quiet',
        action='store_true',
        help='安静模式，减少输出信息'
    )
    
    return parser


def main():
    """
    主函数
    """
    parser = setup_argparse()
    args = parser.parse_args()
    
    # 创建处理器实例
    processor = ImageZipProcessor(verbose=not args.quiet)
    
    try:
        # 处理模式
        if args.get_json:
            result = processor.get_json_info(args.get_json)
            if result:
                print("JSON内容:", json.dumps(result, ensure_ascii=False, indent=2))
            return 0 if result else 1
        
        elif args.get_contents:
            result = processor.get_zip_contents(args.get_contents)
            if result:
                print("ZIP内容:", result)
            return 0 if result else 1
        
        elif args.update_json:
            zip_file, pattern_str = args.update_json
            # 解析pattern（支持逗号分隔的多个值）
            patterns = [p.strip() for p in pattern_str.split(',')]
            success = processor.update_json_info(zip_file, patterns)
            return 0 if success else 1
        
        # 默认处理模式
        elif args.folder:
            if args.dry_run:
                processor.dry_run(args.folder, args.keep_original)
                return 0
            else:
                processed_count = processor.process_folder(args.folder, args.keep_original)
                return 0 if processed_count > 0 else 1
        
        else:
            parser.print_help()
            return 1
            
    except KeyboardInterrupt:
        print("\n用户中断操作")
        return 1
    except Exception as e:
        print(f"处理过程中发生错误: {e}")
        return 1


if __name__ == "__main__":
    exit(main())